import torch
import torch.nn as nn
import numpy as np
from .utils import calc_iou

class FocalCIoULoss(nn.Module):
    def __init__(self, classification_loss_weight=1.0, regression_loss_weight=1.0):
        super(FocalCIoULoss, self).__init__()
        self.classification_loss_weight = classification_loss_weight
        self.regression_loss_weight = regression_loss_weight

    def forward(self, classifications, regressions, anchors, annotation_boxes, annotation_category_ids, **kwargs):
        alpha = 0.25
        gamma = 2.0
        device = classifications.device
        batch_size = classifications.shape[0]
        classification_losses = []
        regression_losses = []

        anchor = anchors[0, :, :]  # assuming all image sizes are the same, which it is
        dtype = anchors.dtype

        anchor_widths = anchor[:, 3] - anchor[:, 1]
        anchor_heights = anchor[:, 2] - anchor[:, 0]
        anchor_ctr_x = anchor[:, 1] + 0.5 * anchor_widths
        anchor_ctr_y = anchor[:, 0] + 0.5 * anchor_heights

        for j in range(batch_size):

            classification = classifications[j, :, :]
            regression = regressions[j, :, :]

            bboxes = annotation_boxes[j]
            category_ids = annotation_category_ids[j]

            classification = torch.clamp(classification, 1e-4, 1.0 - 1e-4)

            if category_ids.shape[0] == 0:
                if torch.cuda.is_available():

                    alpha_factor = torch.ones_like(classification) * alpha
                    alpha_factor = alpha_factor.to(device)
                    alpha_factor = 1. - alpha_factor
                    focal_weight = classification
                    focal_weight = alpha_factor * torch.pow(focal_weight, gamma)

                    bce = -(torch.log(1.0 - classification))

                    cls_loss = focal_weight * bce

                    regression_losses.append(torch.tensor(0).to(dtype).to(device))
                    classification_losses.append(cls_loss.sum())
                else:

                    alpha_factor = torch.ones_like(classification) * alpha
                    alpha_factor = 1. - alpha_factor
                    focal_weight = classification
                    focal_weight = alpha_factor * torch.pow(focal_weight, gamma)

                    bce = -(torch.log(1.0 - classification))

                    cls_loss = focal_weight * bce

                    regression_losses.append(torch.tensor(0).to(dtype))
                    classification_losses.append(cls_loss.sum())

                continue

            IoU = calc_iou(anchor[:, :], bboxes)

            IoU_max, IoU_argmax = torch.max(IoU, dim=1)

            # compute the loss for classification
            targets = torch.ones_like(classification) * -1
            if torch.cuda.is_available():
                targets = targets.to(device)

            targets[torch.lt(IoU_max, 0.4), :] = 0

            positive_indices = torch.ge(IoU_max, 0.5)

            num_positive_anchors = positive_indices.sum()

            assigned_bboxes = bboxes[IoU_argmax, :]
            assigned_category_ids = category_ids[IoU_argmax]

            targets[positive_indices, :] = 0
            targets[positive_indices, assigned_category_ids[positive_indices].long()] = 1

            alpha_factor = torch.ones_like(targets) * alpha
            if torch.cuda.is_available():
                alpha_factor = alpha_factor.to(device)

            alpha_factor = torch.where(torch.eq(targets, 1.), alpha_factor, 1. - alpha_factor)
            focal_weight = torch.where(torch.eq(targets, 1.), 1. - classification, classification)
            focal_weight = alpha_factor * torch.pow(focal_weight, gamma)

            bce = -(targets * torch.log(classification) + (1.0 - targets) * torch.log(1.0 - classification))

            cls_loss = focal_weight * bce

            zeros = torch.zeros_like(cls_loss)
            if torch.cuda.is_available():
                zeros = zeros.to(device)
            cls_loss = torch.where(torch.ne(targets, -1.0), cls_loss, zeros)

            classification_losses.append(cls_loss.sum() / torch.clamp(num_positive_anchors.to(dtype), min=1.0))

            if positive_indices.sum() > 0:
                assigned_bboxes = assigned_bboxes[positive_indices, :]
                assigned_category_ids = assigned_category_ids[positive_indices]

                anchor_widths_pi = anchor_widths[positive_indices]
                anchor_heights_pi = anchor_heights[positive_indices]
                anchor_ctr_x_pi = anchor_ctr_x[positive_indices]
                anchor_ctr_y_pi = anchor_ctr_y[positive_indices]

                regression_pi = regression[positive_indices, :]
                predicted_widths = regression_pi[:, 3].exp() * anchor_widths_pi
                predicted_heights = regression_pi[:, 2].exp() * anchor_heights_pi
                predicted_ctr_x = regression_pi[:, 1] + anchor_ctr_x_pi
                predicted_ctr_y = regression_pi[:, 0] + anchor_ctr_y_pi

                predicted_x_min = predicted_ctr_x - predicted_widths / 2
                predicted_x_max = predicted_ctr_x + predicted_widths / 2
                predicted_y_min = predicted_ctr_y - predicted_heights / 2
                predicted_y_max = predicted_ctr_y + predicted_heights / 2

                predicted_boxes = torch.stack((predicted_x_min, predicted_y_min, predicted_x_max, predicted_y_max)).t()

                predicted_boxes_area = (predicted_boxes[:,3] - predicted_boxes[:,1]) * (predicted_boxes[:,2] - predicted_boxes[:,0])
                assigned_bboxes_area = (assigned_bboxes[:,3] - assigned_bboxes[:,1]) * (assigned_bboxes[:,2] - assigned_bboxes[:,0])
                iw = torch.min(predicted_boxes[:,2], assigned_bboxes[:,2]) - torch.max(predicted_boxes[:,0], assigned_bboxes[:,0])
                ih = torch.min(predicted_boxes[:,3], assigned_bboxes[:,3]) - torch.max(predicted_boxes[:,1], assigned_bboxes[:,1])
                iw = torch.clamp(iw, min=0)
                ih = torch.clamp(ih, min=0)
                intersection = iw * ih
                union = predicted_boxes_area + assigned_bboxes_area - intersection
                IoU = intersection / union

                regression_loss = -torch.log(IoU)

                regression_losses.append(regression_loss.mean())
            else:
                if torch.cuda.is_available():
                    regression_losses.append(torch.tensor(0).to(dtype).to(device))
                else:
                    regression_losses.append(torch.tensor(0).to(dtype))

        return torch.stack(classification_losses).mean(dim=0, keepdim=True) * self.classification_loss_weight, \
               torch.stack(regression_losses).mean(dim=0, keepdim=True) * self.regression_loss_weight
